package com.cmb.lk50.common.receiver;

import io.netty.buffer.UnpooledByteBufAllocator;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.synchronoss.cloud.nio.multipart.util.IOUtils;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;

/**
 * Author       : 蒋俊钊 80256973
 * Date         : 2019/11/19
 * Copyright    (C) ChinaMerchantsBank
 */
public class ResponseBodyReceiver extends ResponseReceiver {

    // 扩展点
    protected String extractBody(String body) {
        return body;
    }

    @Override
    public String receive(HttpServletResponse response) {
        return extractBody(new String(((ResponseBodyReceiverWrapper) response).getBody()));
    }

    public class ResponseBodyReceiverWrapper extends HttpServletResponseWrapper {

        private ByteArrayOutputStream copyStream = new ByteArrayOutputStream();
        private HttpServletResponse response;

        public ResponseBodyReceiverWrapper(HttpServletResponse response) throws IOException {
            super(response);
            this.response = response;
        }

        @Override
        public ServletOutputStream getOutputStream() {
            return new ServletOutputStreamWrapper(copyStream, response);
        }

        @Override
        public PrintWriter getWriter() {
            return new PrintWriter(copyStream);
        }

        public byte[] getBody() {
            return copyStream.toByteArray();
        }


        private class ServletOutputStreamWrapper extends ServletOutputStream {

            private ByteArrayOutputStream byteArrayOutputStream;
            private HttpServletResponse response;

            public ServletOutputStreamWrapper(ByteArrayOutputStream byteArrayOutputStream, HttpServletResponse response) {
                this.byteArrayOutputStream = byteArrayOutputStream;
                this.response = response;
            }

            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public void setWriteListener(WriteListener listener) {
            }

            @Override
            public void write(int b) throws IOException {
                this.byteArrayOutputStream.write(b);
            }

            @Override
            public void write(byte[] b) throws IOException {
                this.byteArrayOutputStream.write(b);
            }

            @Override
            public void write(byte[] b, int off, int len) throws IOException {
                this.byteArrayOutputStream.write(b, off, len);
            }

            @Override
            public void flush() throws IOException {
                if (!this.response.isCommitted()) {
                    byte[] body = this.byteArrayOutputStream.toByteArray();
                    ServletOutputStream outputStream = this.response.getOutputStream();
                    outputStream.write(body);
                    outputStream.flush();
                }
            }
        }
    }

    //////////////////////////////////////////////////////////////////////////

    /**
     * Webflux下的body提取
     */
    @Override
    public String receive(ServerHttpResponse response) {
        if (response instanceof WrapperServerHttpResponseDecorator) {
            WrapperServerHttpResponseDecorator decorator = (WrapperServerHttpResponseDecorator) response;
            return extractBody(decorator.getBodyStr());
        } else {
            return extractBody("error");
        }
    }

    public class WrapperServerHttpResponseDecorator extends ServerHttpResponseDecorator {

        private String bodyStr;

        public WrapperServerHttpResponseDecorator(ServerHttpResponse delegate) {
            super(delegate);
        }

        public String getBodyStr() {
            return bodyStr;
        }

        @Override
        public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
            return setComplete();
        }

        @Override
        public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
            if (body instanceof Mono) {
                final Mono<DataBuffer> monoBody = (Mono<DataBuffer>) body;
                return super.writeWith(monoBody.map(this::retain));
            } else if (body instanceof Flux) {
                final Flux<DataBuffer> fluxBody = (Flux<DataBuffer>) body;
                return super.writeWith(fluxBody.map(this::retain));
            }
            return super.writeWith(body);
        }

        private <T extends DataBuffer> T retain(DataBuffer buffer) {
            try {
                InputStream inputStream = buffer.asInputStream();
                bodyStr = IOUtils.inputStreamAsString(inputStream, "UTF-8");
                NettyDataBufferFactory factory = new NettyDataBufferFactory(new UnpooledByteBufAllocator(false));
                return (T) factory.wrap(bodyStr.getBytes());
            } catch (IOException e) {
                e.printStackTrace();
            } finally {
                DataBufferUtils.release(buffer);
            }
            return null;
        }

    }

}
